'''
author:        Wang Chenyang <cy-wang21@mails.tsinghua.edu.cn>
date:          2024-09-23
Copyright © Department of Physics, Tsinghua University. All rights reserved
'''

import requests
from xml.etree import ElementTree as ET


def get_bibtex_for_multiple_papers(arxiv_ids):
    # 将 arXiv IDs 转换为逗号分隔的字符串
    id_list = ','.join(arxiv_ids)

    # 构建API请求URL
    url = f"https://export.arxiv.org/api/query?id_list={id_list}"

    # 发送请求
    response = requests.get(url)
    # with open("data/temp.xml", "rb") as fp:
    #     response_content = fp.read()

    with open("data/temp.xml", "wb") as fp:
        fp.write(response.content)
    return

    root = ET.fromstring(response_content)

    NS = {'atom': 'http://www.w3.org/2005/Atom',
          'dc': 'http://purl.org/dc/elements/1.1/'}

    # 提取BibTeX引用
    bibtex_dict = {}
    for entry in root.findall('atom:entry', namespaces=NS):
        arxiv_id = entry.find('atom:id', namespaces=NS).text.split('/')[-1]
        print(arxiv_id)
        for link in entry.findall('.//atom:link', namespaces=NS):
            if 'title' in link.attrib and link.attrib['title'] == 'bibtex':
                bibtex_url = link.attrib['href']
                bibtex_response = requests.get(bibtex_url)
                bibtex_dict[arxiv_id] = bibtex_response.text.strip()
                break

    return bibtex_dict


# 使用示例
arxiv_ids = ['1404.7739', '1609.04468', '1907.10599v4']
bibtex_references = get_bibtex_for_multiple_papers(arxiv_ids)
print(bibtex_references)

for arxiv_id, bibtex in bibtex_references.items():
    print(f"BibTeX for {arxiv_id}:")
    print(bibtex)
    print("\n" + "-"*50 + "\n")
